import numpy as np
from lazypredict.Supervised import LazyClassifier
import logging
from util import createXY
import faiss
import pickle
from sklearn.model_selection import train_test_split

# 配置logging, 确保能够打印正在运行的函数名
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 载入和预处理数据
X, y = createXY(train_folder="../data/train", dest_folder=".")
if X is None or y is None:
    logging.error("createXY函数返回了None，程序将退出。")
    exit(1)

X = np.array(X).astype('float32')
faiss.normalize_L2(X)  # 对数据进行L2归一化
y = np.array(y)
logging.info("数据加载和预处理完成。")

# 数据集分割为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
logging.info("数据集划分为训练集和测试集。")

clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
result, _ = clf.fit(X_train, X_test, y_train, y_test)
print(result)

# 获取F1分数最高的模型
if 'F1 Score' in result:
    best_model_name = result['F1 Score'].idxmax()  # 获取F1分数最高行的索引值，即：模型名称
    print("\nF1分数最高的模型是:", best_model_name)
else:
    logging.error("结果中没有'F1 Score'列。")
    exit(1)

# clf.models 是包含所有训练过的模型 (名称, 模型对象) 键值对的字典
best_model = clf.models[best_model_name]  # 根据模型名称，从模型字典中获取模型对象

result = best_model.predict(X_test)  # 该字典可以直接被拿来进行预测
print(f"用{best_model_name}预测X_test的结果是:\n{result}")
pickle.dump(best_model, open('best_model.pkl', 'wb'))